import torch
from PIL import Image
import torchvision.transforms as tfs
import matplotlib.pyplot as plt


class ImgFactory(object):
    def __init__(self, patch=2):
        super(ImgFactory, self).__init__()
        self.patch = patch
        self.im_tfs = tfs.Compose([
            tfs.ToTensor(),
            tfs.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def getImagePatch(self, filename):
        img = Image.open(filename)
        width, height = img.size
        num_patch_w = width // self.patch
        num_patch_h = height // self.patch

        patch_list = []

        num = 1
        for i in range(num_patch_h):
            for j in range(num_patch_w):
                s_y = i * self.patch
                s_x = j * self.patch
                box = (s_x, s_y, self.patch + s_x, self.patch + s_y)
                region = img.crop(box)
                patch_list.append(region)
                plt.subplot(num_patch_h, num_patch_w, num), plt.imshow(region), plt.axis("off")
                num = num + 1

        plt.savefig("patch.png")

        for i in range(len(patch_list)):
            patch_list[i] = self.im_tfs(patch_list[i])
            patch_list[i] = patch_list[i].view(1, -1)

        seq = torch.cat(patch_list, dim=0)
        return seq


if __name__ == "__main__":
    factory = ImgFactory()
    seq = factory.getImagePatch(r"D:\Saliency\result.png")
    print(seq.shape)